package com.gitee.wsl.func.bijection.ext

import com.gitee.wsl.func.bijection.Bijection
import com.gitee.wsl.func.Function

object Bijections {
    val _identity: Bijection<Any, Any> = IdentityBijection()

    fun <A> identity(): Bijection<A, A> {
        return _identity as Bijection<A, A>
    }

    /**
     * 将两个函数配对形成一个双射
     *
     * @param f1
     * @param f2
     * @return
     */
    fun <A, B> pair(
        f1: Function<A, B>,
        f2: Function<B, A>,
    ): Bijection<A, B> {
        return PairBijection(f1, f2)
    }

    /**
     * 将一个双射反向
     *
     * @param bij
     * @return
     */
    fun <B, A> invert(bij: Bijection<A, B>): Bijection<B, A> {
        if (bij is InvertBijection<*, *>) {
            return bij.original as Bijection<B, A>
        }
        return InvertBijection(bij)
    }

    /**
     * 将两个双射组合成一个新的双射。正向映射时先调用b1, 再调用b2, 逆向映射时先调用b2再调用b1
     *
     * @param b1
     * @param b2
     * @return
     */
    fun <A, B, C> compose(b1: Bijection<A, B>, b2: Bijection<B, C>): Bijection<A, C> {
        return ComposeBijection(b1, b2)
    }

    fun <T, R> adapt(
        fn: Function<R, R>,
        b: Bijection<T, R>,
    ): Function<T, T> {
        return AdaptFunction<T, R>(b, fn)
    }

    internal class AdaptFunction<T, R>(
        val b: Bijection<T, R>,
        val fn: Function<R, R>
    ) : Function<T, T> {

        override fun apply(t: T): T {
            var r: R = b.apply(t)
            r = fn.apply(r)
            return b.invert(r)
        }
    }

    internal class PairBijection<A, B>(
        val f1: Function<A, B>,
        val f2: Function<B, A>,
    ) : Bijection<A, B> {

        public override fun apply(a: A): B {
            return f1.apply(a)
        }

        public override fun invert(b: B): A {
            return f2.apply(b)
        }
    }

    internal class IdentityBijection : Bijection<Any, Any> {
        public override fun apply(a: Any): Any {
            return a
        }

        public override fun invert(b: Any): Any {
            return b
        }
    }

    internal class InvertBijection<B, A>(private val bij: Bijection<A, B>) : Bijection<B, A> {

        val original: Bijection<A, B>
            get() = bij

        public override fun apply(b: B): A {
            return bij.invert(b)
        }

        public override fun invert(a: A): B {
            return bij.apply(a)
        }
    }

    internal class ComposeBijection<A, B, C>(
        private val b1: Bijection<A, B>,
        private val b2: Bijection<B, C>
    ) : Bijection<A, C> {

        public override fun apply(a: A): C {
            val b: B = b1.apply(a)
            return b2.apply(b)
        }

        public override fun invert(c: C): A {
            val b: B = b2.invert(c)
            return b1.invert(b)
        }
    }
}